package com.yifeng.repo.controller.kafka.message.consume;

import com.google.common.base.Strings;
import com.yifeng.repo.base.constant.BaseConstant;
import com.yifeng.repo.controller.kafka.message.constant.KafkaConstant;
import lombok.extern.slf4j.Slf4j;
import org.apache.kafka.clients.consumer.ConsumerInterceptor;
import org.apache.kafka.clients.consumer.ConsumerRecord;
import org.apache.kafka.clients.consumer.ConsumerRecords;
import org.apache.kafka.clients.consumer.OffsetAndMetadata;
import org.apache.kafka.common.TopicPartition;
import org.apache.kafka.common.header.Header;

import java.util.*;
import java.util.stream.Collectors;

import static com.yifeng.repo.controller.kafka.message.constant.KafkaConstant.SUB_EXPRESSION_ALL;

/**
 * Created by daibing on 2024/5/2.
 */
@Slf4j
public class TagConsumerInterceptor implements ConsumerInterceptor<String, String> {

    /**
     * 支持表达式
     */
    private boolean subscribeAll;
    private Set<String> subExpressionSet;

    /**
     * 在消息被消费前执行逻辑
     * 该方法在消息返回给 Consumer 程序之前调用。
     * 也就是说在开始正式处理消息之前，拦截器会先拦一道，搞一些事情，之后再返回给你。
     */
    @Override
    public ConsumerRecords<String, String> onConsume(ConsumerRecords<String, String> records) {
        // 空集合直接返回
        if (records.isEmpty()) {
            return records;
        }

        // 如果订阅全部，就直接返回。
        if (subscribeAll) {
            return records;
        }

        // 多个tag是 TagA || TagC || TagD
        Map<TopicPartition, List<ConsumerRecord<String, String>>> partition2Records = new HashMap<>(records.partitions().size());
        records.partitions().forEach(partition -> {
            records.records(partition).forEach(record -> {
                Header tagHeader = record.headers().lastHeader(KafkaConstant.PROPERTY_TAGS);
                String tag = Optional.ofNullable(tagHeader).map(s -> new String(tagHeader.value()).trim()).orElse(BaseConstant.EMPTY);
                if (Strings.isNullOrEmpty(tag) || subExpressionSet.contains(tag)) {
                    partition2Records.computeIfAbsent(partition, k -> new ArrayList<>());
                    partition2Records.get(partition).add(record);
                }
            });
        });
        return new ConsumerRecords<>(partition2Records);
    }

    /**
     * 在消费者提交偏移量时执行逻辑
     * Consumer 在提交位移之后调用该方法。
     * 通常你可以在该方法中做一些记账类的动作，比如打日志等。
     */
    @Override
    public void onCommit(Map<TopicPartition, OffsetAndMetadata> offsets) {

        // log.info("offsets: {}", offsets);
    }

    /**
     * 在拦截器关闭时执行清理逻辑
     */
    @Override
    public void close() {

    }

    /**
     * 获取配置信息
     */
    @Override
    public void configure(Map<String, ?> configs) {
        // 未指定订阅表达式就默认全部是 *
        Object obj = configs.get(KafkaConstant.PROPERTY_TAGS);
        if (obj == null) {
            this.subscribeAll = true;
            return;
        }

        // 直接指定是全部 *
        String str = String.valueOf(obj).trim();
        if (SUB_EXPRESSION_ALL.equals(str)) {
            this.subscribeAll = true;
            return;
        }

        // 指定具体表达式，支持多个tag TagA || TagC || TagD
        this.subscribeAll = false;
        String[] split = str.split(KafkaConstant.SUB_EXPRESSION_SEPARATOR);
        this.subExpressionSet = Arrays.stream(split).map(String::trim).collect(Collectors.toSet());
    }
}
